package handler

import (
	"context"
	"fmt"
	"github.com/cloudwego/hertz/pkg/app"
	"github.com/hertz-contrib/websocket"
	"strconv"
	"strings"
	"time"
	"view-crawler-backend/internal/constant"
	"view-crawler-backend/internal/dao"
	"view-crawler-backend/internal/service"
	"view-crawler-backend/internal/tools"
)

// 用于将连接升级成WebSocket连接
var upgrader = websocket.HertzUpgrader{
	CheckOrigin: func(ctx *app.RequestContext) bool {
		return true
	},
}

// TaskExec 执行任务，并基于WebSocket实时交互
func TaskExec(_ context.Context, c *app.RequestContext) {
	upgrader.Upgrade(c, func(conn *websocket.Conn) {
		msgRcvStream := make(chan string, 5)
		msgSendStream := make(chan string, 5)

		dataStream := make(chan map[string]any, 5)
		done := make(chan struct{})
		controlStream := make(chan bool)
		// 用于向协调协程反馈WebSocket连接关闭，或是因为出错需要关闭连接
		connDone := make(chan bool, 5)

		defer close(done)

		// 启动指令接收协程
		go messageReceiveGoroutine(conn, msgRcvStream, connDone)
		// 开始协调各协程
		for {
			select {
			case message := <-msgRcvStream:
				tp, ct, err := parseMessage(message)
				if err != nil {
					msgSendStream <- ""
					return
				}
				switch tp {
				case constant.START:
					taskId, err := strconv.ParseInt(ct, 10, 64)
					if err != nil {
						msgSendStream <- ""
						return
					}
					// 启动爬虫协程
					go func() {
						service.TaskExec(taskId, msgSendStream, dataStream, done, connDone, controlStream)
						// 任务执行完成，告知前端，前端可以发起连接关闭
						connDone <- false
						return
					}()
					// 启动消息发送协程（主要是爬取过程的日志）
					go messageSendGoroutine(conn, msgSendStream, done, connDone)
					// 启动数据持久化协程
					go dataPersistentGoroutine(dataStream, done, connDone, taskId)
				case constant.PAUSE:
					controlStream <- false
				case constant.CONTINUE:
					controlStream <- true
				case constant.END:
					msgSendStream <- tools.LogFormat(constant.RED, constant.TASK_END, "用户终止了任务")
					connDone <- false
				}
			case flag := <-connDone:
				// 由于发生了错误或任务完成或用户终止了任务，需要后端发起WebSocket连接断开
				if !flag {
					// 自旋等待，将残留日志发送
					for len(msgSendStream) > 0 {
					}
					// 等待一段时间，确保日志被正常接收
					time.Sleep(50 * time.Millisecond)
					msgSendStream <- ""
				}
				return
			}
		}
	})
}

func parseMessage(message string) (int, string, error) {
	arr := strings.SplitN(message, " ", 2)
	param2 := arr[1]
	param1, err := strconv.ParseInt(arr[0], 10, 64)
	if err != nil {
		return 0, "", err
	}
	return int(param1), param2, nil
}

func messageReceiveGoroutine(conn *websocket.Conn, msgRcvStream chan string, connDone chan bool) {
	for {
		_, message, err := conn.ReadMessage()
		// 连接异常
		if err != nil {
			connDone <- true
			return
		}
		msgRcvStream <- string(message)
	}
}

func messageSendGoroutine(conn *websocket.Conn, msgSendStream chan string, done chan struct{}, connDone chan bool) {
	for {
		select {
		case message := <-msgSendStream:
			var err error
			if message == "" {
				err = conn.WriteMessage(websocket.CloseMessage,
					websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
				return
			} else {
				err = conn.WriteMessage(websocket.TextMessage, []byte(message))
				fmt.Println(message)
			}
			if err != nil {
				connDone <- true
			}
		case <-done:
			return
		}
	}
}

func dataPersistentGoroutine(dataStream chan map[string]any, done chan struct{}, connDone chan bool, taskId int64) {
	execId, err := dao.QueryTaskExecTimesById(dao.DB, taskId)
	if err != nil {
		connDone <- false
		return
	}

	// 每10条数据写一次DB
	dataList, idx := make([]map[string]any, 10), 0
	for {
		select {
		case data := <-dataStream:
			if idx == 10 {
				err = dao.BatchAddExecData(dao.DB, taskId, dataList)
				if err != nil {
					connDone <- false
					return
				}
				idx = 0
			} else {
				data["execId"] = execId
				dataList[idx] = data
				idx++
			}
		case <-done:
			// 最后一批数据
			if idx > 0 {
				_ = dao.BatchAddExecData(dao.DB, taskId, dataList[:idx])
			}
			return
		}
	}
}
